import json
import os
import signal
import threading
import time
from urllib.parse import unquote
from urllib.parse import urlparse

import requests
from mitmproxy import ctx

from m3u8_downloader import M3u8Downloader


class MediaSaver:
    # 类变量，控制是否保存媒体文件
    save_enabled = False
    toggle_count = 0

    @classmethod
    def signal_handler(cls, signum, frame):
        """处理Ctrl+C信号，切换保存模式"""
        cls.toggle_count += 1
        cls.save_enabled = not cls.save_enabled
        status = "开启" if cls.save_enabled else "关闭"
        ctx.log.info(f"\n已{status}媒体保存模式")
        if cls.toggle_count >= 2:
            # 重置计数
            cls.toggle_count = 0

    def __init__(self):
        # 支持的媒体文件扩展名
        self.media_extensions = [".mp4", ".mp3", ".m3u8", ".ts"]
        # 保存的基础目录
        self.save_dir = os.path.expanduser("~/Downloads/media_saver")
        # 创建保存目录
        os.makedirs(self.save_dir, exist_ok=True)
        # 记录正在下载的m3u8文件
        self.downloading_m3u8 = {}  # 改为字典，存储m3u8相关信息
        # m3u8下载器
        self.m3u8_downloader = M3u8Downloader()
        # ts文件缓存，存储文件路径而不是内容
        self.ts_cache = {}
        # 设置信号处理器
        signal.signal(signal.SIGINT, self.signal_handler)

    def get_filename_from_url(self, url):
        """从URL中提取文件名"""
        # 解码URL
        url = unquote(url)
        # 解析URL
        parsed = urlparse(url)
        # 获取路径部分
        path = parsed.path
        # 提取文件名
        filename = os.path.basename(path)
        # 如果文件名为空，使用时间戳
        if not filename:
            filename = f"media_{int(time.time())}"
        # 如果没有扩展名，根据Content-Type添加
        if not os.path.splitext(filename)[1]:
            filename += ".mp4"
        return filename

    def is_media_file(self, url, content_type):
        """判断是否是媒体文件"""
        # 检查URL扩展名
        if any(ext in url.lower() for ext in self.media_extensions):
            return True

        # 检查Content-Type
        media_types = [
            "video/",
            "audio/",
            "application/x-mpegURL",
            "application/vnd.apple.mpegurl",
            "application/octet-stream",
        ]
        return any(t in content_type.lower() for t in media_types)

    def download_file(self, url, name):
        path = os.path.join(self.save_dir, name)
        if os.path.exists(path):
            print(f"#################### {name} 已存在 ####################")
            return
        print(f"保存到: {path}")
        try:
            headers = {
                "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
            }
            res = requests.get(
                url, stream=True, headers=headers, timeout=10, verify=False
            )
            print(f"#################### res.status_code: {res.status_code} ####################")
            if res.status_code == 200:
                with open(path, "wb") as f:
                    for chunk in res.iter_content(chunk_size=1024):
                        f.write(chunk)
                ctx.log.info(f"保存文件: {name}")
            else:
                ctx.log.error(f"保存文件失败: {name}")
        except Exception as e:
            ctx.log.error(f"保存文件失败: {name} {e}")

    def parse_audio(self, content):
        print(f"#################### {content} ####################")
        data = json.loads(content).get("data")
        if data:
            resourceDTO = data.get("resourceDTO")
            downUrl = resourceDTO.get("downUrl")
            title = resourceDTO.get("title")
            name = f"{title}.mp3"
            print(f"#################### {title} ####################")
            print(f"#################### {downUrl} ####################")
            t = threading.Thread(target=self.download_file, args=(downUrl, name))
            t.start()

    def response(self, flow):
        """处理响应"""
        # if "aiaudio.do" in flow.request.url:
        #     self.parse_audio(flow.response.content)
        # return

        try:
            # 获取响应内容类型
            content_type = flow.response.headers.get("content-type", "")

            # 检查是否是媒体文件
            if not self.is_media_file(flow.request.url, content_type):
                return

            # 如果保存模式未启用，直接返回
            if not self.save_enabled:
                return

            # 获取文件名
            filename = self.get_filename_from_url(flow.request.url)

            # 创建保存路径
            save_path = os.path.join(self.save_dir, filename)
            if os.path.exists(save_path):
                name, ext = os.path.splitext(save_path)
                save_path = name + "_" + int(time.time()) + ext

            # 如果是m3u8文件
            if filename.endswith(".m3u8"):
                # 避免重复下载
                if flow.request.url in self.downloading_m3u8:
                    return

                # 保存m3u8文件
                with open(save_path, "wb") as f:
                    f.write(flow.response.content)

                # 创建临时目录用于存储ts文件
                ts_dir = os.path.splitext(save_path)[0] + "_ts"
                os.makedirs(ts_dir, exist_ok=True)

                # 记录m3u8信息
                self.downloading_m3u8[flow.request.url] = {
                    "save_path": save_path,
                    "output_file": os.path.splitext(save_path)[0] + ".mp4",
                    "base_url": os.path.dirname(flow.request.url) + "/",
                    "ts_files": set(),  # 记录需要的ts文件
                    "ts_dir": ts_dir,  # 添加ts文件保存目录
                }

                # 解析m3u8文件，获取所需的ts文件列表
                self.m3u8_downloader.parse_m3u8(
                    flow.request.url,
                    save_path,
                    self.downloading_m3u8[flow.request.url]["ts_files"],
                )

                ctx.log.info(f"开始处理m3u8: {filename}")

            # 如果是ts文件
            elif filename.endswith(".ts"):
                # 检查是否是某个m3u8的一部分
                for m3u8_info in self.downloading_m3u8.values():
                    if flow.request.url in m3u8_info["ts_files"]:
                        # 获取ts文件序号
                        ts_index = len(self.ts_cache.get(flow.request.url, []))
                        # 构建ts文件保存路径
                        ts_path = os.path.join(
                            m3u8_info["ts_dir"], f"{ts_index:05d}.ts"
                        )
                        # 保存ts文件
                        with open(ts_path, "wb") as f:
                            f.write(flow.response.content)
                        # 缓存ts文件路径
                        self.ts_cache[flow.request.url] = ts_path
                        # 检查是否所有ts文件都已获取
                        self.check_and_merge_m3u8(flow.request.url)
                        return

                # 如果不是m3u8的一部分，直接保存
                with open(save_path, "wb") as f:
                    f.write(flow.response.content)
                ctx.log.info(f"已保存: {filename}")

            # 如果是其他媒体文件
            else:
                # 保存文件
                with open(save_path, "wb") as f:
                    f.write(flow.response.content)
                ctx.log.info(f"已保存: {filename}")

        except Exception as e:
            ctx.log.error(f"保存文件失败: {str(e)}")

    def check_and_merge_m3u8(self, ts_url):
        """检查并尝试合并m3u8文件"""
        try:
            # 查找这个ts文件属于哪个m3u8
            for m3u8_url, m3u8_info in self.downloading_m3u8.items():
                if ts_url in m3u8_info["ts_files"]:
                    # 检查是否所有ts文件都已获取
                    all_ts_urls = m3u8_info["ts_files"]
                    cached_ts_urls = set(self.ts_cache.keys())
                    if all_ts_urls.issubset(cached_ts_urls):
                        ctx.log.info("所有ts文件已获取，开始合并...")
                        # 开始合并
                        self.m3u8_downloader.merge_from_cache(
                            m3u8_info["save_path"],
                            m3u8_info["output_file"],
                            m3u8_info["ts_dir"],
                        )
                        # 清理缓存
                        for url in all_ts_urls:
                            self.ts_cache.pop(url, None)
                        self.downloading_m3u8.pop(m3u8_url, None)
                        ctx.log.info(f"合并完成: {m3u8_info['output_file']}")
                    break
        except Exception as e:
            ctx.log.error(f"合并失败: {str(e)}")


addons = [MediaSaver()]
